package com.yoyosys.db.partition.query.impl;

import com.yoyosys.db.partition.constant.PartitionType;
import com.yoyosys.db.partition.query.PartitionQuery;
import com.yoyosys.db.partition.vo.Partition;
import com.yoyosys.db.partition.vo.PartitionInfo;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;

/**
 * PG 分区分桶信息查询器
 */
public class PostgreSqlPartitionQuery implements PartitionQuery {
    private static PostgreSqlPartitionQuery instance = new PostgreSqlPartitionQuery();
    private PostgreSqlPartitionQuery(){}
    public static PostgreSqlPartitionQuery getInstance(){
        return instance;
    }

    /**
     * 查询表分区信息
     *
     * @param conn
     * @param schemaName 库名
     * @param tableName
     * @return
     */
    @Override
    public PartitionInfo getPartition(Connection conn, String schemaName, String tableName) throws SQLException {
        PartitionInfo info = new PartitionInfo();
        info.setSchemaName(schemaName);
        info.setTableName(tableName);
        info.setPartitions(new ArrayList<>());

        // 第1步：检查表是否为分区表（主表）
        String checkPartitionedSql = new StringBuilder()
                .append("SELECT pt.partrelid IS NOT NULL AS is_partitioned ")
                .append("FROM pg_class c ")
                .append("JOIN pg_namespace ns ON c.relnamespace = ns.oid ")
                .append("LEFT JOIN pg_partitioned_table pt ON c.oid = pt.partrelid ")
                .append("WHERE ns.nspname = ? AND c.relname = ?")
                .toString();
        System.out.println("checkPartitionedSql: " + checkPartitionedSql);
        boolean isPartitioned = false;
        try (PreparedStatement pstmt = conn.prepareStatement(checkPartitionedSql)) {
            pstmt.setString(1, schemaName);
            pstmt.setString(2, tableName);
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    isPartitioned = rs.getBoolean("is_partitioned");
                }
            }
        }

        if (!isPartitioned) {
            info.setPartitionType(PartitionType.NONE.name());
            return info;
        }

        // 第2步：获取分区类型和键
        String partitionTypeKeySql = new StringBuilder()
                .append(" SELECT  ")
                .append("     a.attname AS partition_key, ")
                .append("     pg_get_expr(p.partexprs, p.partrelid) AS partition_expression, ")
                .append("     p.partstrat, ")
                .append("     p.partnatts AS number_of_partition_columns ")
                .append(" FROM  ")
                .append("     pg_partitioned_table p ")
                .append(" JOIN  ")
                .append("     pg_attribute a ON p.partrelid = a.attrelid AND a.attnum = ANY(p.partattrs) ")
                .append(" WHERE  ")
                .append("     p.partrelid = ?::regclass")
                .toString();
        System.out.println("partitionTypeKeySql: " + partitionTypeKeySql);
        try (PreparedStatement pstmt = conn.prepareStatement(partitionTypeKeySql)) {
            pstmt.setString(1, schemaName + "." + tableName);
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    info.setPartitionKey(rs.getString("partition_key"));

                    char strat = rs.getString("partstrat").charAt(0);
                    String type = "";
                    switch (strat) {
                        case 'r': type = PartitionType.RANGE.name(); break;
                        case 'l': type = PartitionType.LIST.name(); break;
                        case 'h': type = PartitionType.HASH.name(); break;
                        default: type = PartitionType.UNKNOWN.name();
                    }
                    info.setPartitionType(type);
                }
            }
        }


        // 第4步：获取所有子分区
        String partitionsSql = new StringBuilder()
                .append("SELECT c.relname AS partition_name ")
                .append("FROM pg_inherits i ")
                .append("JOIN pg_class parent ON i.inhparent = parent.oid ")
                .append("JOIN pg_class c ON i.inhrelid = c.oid ")
                .append("JOIN pg_namespace ns ON parent.relnamespace = ns.oid ")
                .append("WHERE ns.nspname = ? AND parent.relname = ?")
                .toString();
        System.out.println("partitionsSql: " + partitionsSql);
        try (PreparedStatement pstmt = conn.prepareStatement(partitionsSql)) {
            pstmt.setString(1, schemaName);
            pstmt.setString(2, tableName);
            try (ResultSet rs = pstmt.executeQuery()) {
                while (rs.next()) {
                    String partitionName = rs.getString("partition_name");
                    Partition partition = new Partition();
                    partition.setName(partitionName);

                    // 获取分区约束表达式
                    String constraintSql = new StringBuilder()
                            .append("SELECT pg_get_expr(conbin, conrelid) AS constraint_expr ")
                            .append("FROM pg_constraint ")
                            .append("WHERE conrelid = ?::regclass AND contype = 'c'")
                            .toString();
                    System.out.println("constraintSql: " + constraintSql);
                    try (PreparedStatement constraintPstmt = conn.prepareStatement(constraintSql)) {
                        constraintPstmt.setString(1, schemaName + "." + partitionName);
                        try (ResultSet constraintRs = constraintPstmt.executeQuery()) {
                            if (constraintRs.next()) {
                                partition.setExpression(constraintRs.getString("constraint_expr"));
                            }
                        }
                    }

                    // 获取分区行数
                    //取估计值
                    /*
                    String rowCountSql = new StringBuilder()
                            .append("SELECT reltuples::bigint AS row_count ") //reltuples是估计值，可能不是精确计数
                            .append("FROM pg_class c ")
                            .append("JOIN pg_namespace ns ON c.relnamespace = ns.oid ")
                            .append("WHERE ns.nspname = ? AND c.relname = ?")
                            .toString();
                    */

                    //取精确值
                    String rowCountSql = new StringBuilder()
                            .append(" SELECT ")
                            .append("     child.relname AS partition_name, ")
                            .append("     (xpath('/row/cnt/text()', query_to_xml(format('SELECT count(*) AS cnt FROM %I.%I',  ")
                            .append("                                         child.relnamespace::regnamespace::text,  ")
                            .append("                                         child.relname),  ")
                            .append("                                 false,  ")
                            .append("                                 true,  ")
                            .append("                                 '')))[1]::text::int AS row_count ")
                            .append(" FROM pg_inherits ")
                            .append(" JOIN pg_class parent ON pg_inherits.inhparent = parent.oid ")
                            .append(" JOIN pg_class child ON pg_inherits.inhrelid = child.oid ")
                            .append(" JOIN pg_namespace nmsp ON nmsp.oid = parent.relnamespace ")
                            .append(" WHERE 1=1 ")
                            .append("     AND nmsp.nspname = ? ")
                            .append("     AND parent.relname = ? ")
                            .append("     AND child.relname = ? ")
                            .toString();
                    System.out.println("rowCountSql: " + rowCountSql);
                    try (PreparedStatement rowCountPstmt = conn.prepareStatement(rowCountSql)) {
                        rowCountPstmt.setString(1, schemaName);
                        rowCountPstmt.setString(2, tableName);
                        rowCountPstmt.setString(3, partitionName);
                        try (ResultSet rowCountRs = rowCountPstmt.executeQuery()) {
                            if (rowCountRs.next()) {
                                partition.setRowCount(rowCountRs.getLong("row_count"));
                            }
                        }
                    }

                    info.getPartitions().add(partition);
                }
            }
        }

        return info;
    }

}
